import cv2
import numpy as np
import threading
import time
from pyscrcpy import Client
import subprocess
import os
import json
import argparse
import sys
from pathlib import Path
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout,
                             QHBoxLayout, QWidget, QLabel, QPushButton,
                             QStatusBar, QFrame, QMenuBar, QAction, QFileDialog, QMessageBox)
from PyQt5.QtCore import Qt, QTimer, pyqtSignal, QThread, QRect
from PyQt5.QtGui import QImage, QPixmap, QPainter, QPen, QFont

class VideoWidget(QLabel):
    """自定义视频显示控件"""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.parent_window = parent
        self.setMinimumSize(800, 600)
        self.setStyleSheet("border: 1px solid gray;")
        self.setAlignment(Qt.AlignCenter)
        self.setScaledContents(False)

        # 绘制相关
        self.drawing = False
        self.start_point = None
        self.current_rect = None

        # 显示参数
        self.video_rect = QRect()  # 视频在控件中的实际显示区域
        self.scale_factor = 1.0

    def set_video_rect(self, rect, scale_factor):
        """设置视频显示区域"""
        self.video_rect = rect
        self.scale_factor = scale_factor

    def mousePressEvent(self, event):
        """鼠标按下事件"""
        if not self.video_rect.contains(event.pos()):
            return

        if event.button() == Qt.LeftButton:
            self.drawing = True
            self.start_point = event.pos()
            self.current_rect = None

        elif event.button() == Qt.RightButton:
            # 右键删除选框
            if self.parent_window:
                self.parent_window.delete_roi_at_position(event.pos())

    def mouseMoveEvent(self, event):
        """鼠标移动事件"""
        if self.drawing and self.start_point and self.video_rect.contains(event.pos()):
            # 创建当前绘制的矩形
            x1, y1 = self.start_point.x(), self.start_point.y()
            x2, y2 = event.pos().x(), event.pos().y()

            self.current_rect = QRect(
                min(x1, x2), min(y1, y2),
                abs(x2 - x1), abs(y2 - y1)
            )
            self.update()  # 触发重绘

    def mouseReleaseEvent(self, event):
        """鼠标释放事件"""
        if self.drawing and self.start_point and event.button() == Qt.LeftButton:
            if self.current_rect and self.current_rect.width() > 10 and self.current_rect.height() > 10:
                # 添加新的ROI
                if self.parent_window:
                    self.parent_window.add_roi(self.current_rect)

            self.drawing = False
            self.start_point = None
            self.current_rect = None
            self.update()

    def paintEvent(self, event):
        """重绘事件"""
        super().paintEvent(event)

        if not self.parent_window:
            return

        painter = QPainter(self)

        # 绘制已确定的ROI
        pen = QPen(Qt.green, 2)
        painter.setPen(pen)
        font = QFont()
        font.setPointSize(10)
        painter.setFont(font)

        for i, roi_rect in enumerate(self.parent_window.display_rois):
            painter.drawRect(roi_rect)
            painter.drawText(roi_rect.x(), roi_rect.y() - 5, f"ROI-{i+1}")

        # 绘制正在创建的ROI
        if self.current_rect:
            pen = QPen(Qt.red, 2)
            painter.setPen(pen)
            painter.drawRect(self.current_rect)

class AndroidRelocatorGUI(QMainWindow):
    """主窗口类"""

    def __init__(self, device_ip=None, config_path=None, aspect_ratio=1.8526, max_size=1778):
        super().__init__()

        # 初始化参数
        self.device_ip = device_ip or self.auto_detect_device()
        # 如果没有指定配置路径，使用默认的AppData路径
        self.config_path = config_path or self.get_default_config_path()
        self.aspect_ratio = aspect_ratio
        self.max_size = max_size

        # 数据存储
        self.normalized_rois = []  # 归一化的ROI坐标
        self.display_rois = []     # 显示坐标的ROI
        self.latest_frame = None
        self.frame_lock = threading.Lock()

        # 连接状态
        self.client = None
        self.connected = False
        self.reconnect_attempts = 0
        self.max_reconnect_attempts = 3

        # 视频参数
        self.frame_width = 0
        self.frame_height = 0

        # 初始化界面
        self.init_ui()
        
        # 更新窗口标题显示配置文件路径
        self.update_window_title()

        # 加载配置
        self.load_config()

        # 设置定时器更新显示
        self.update_timer = QTimer()
        self.update_timer.timeout.connect(self.update_display)
        self.update_timer.start(33)  # 约30fps

        # 连接设备
        self.connect_client()

    def get_default_config_path(self):
        """获取默认配置文件路径"""
        # 获取用户的AppData\Local目录
        if os.name == 'nt':  # Windows系统
            appdata_local = os.path.join(os.environ['LOCALAPPDATA'], 'AndroidRelocator')
        else:  # Linux/Mac系统
            appdata_local = os.path.join(os.path.expanduser('~'), '.config', 'AndroidRelocator')
        
        # 确保目录存在
        Path(appdata_local).mkdir(parents=True, exist_ok=True)
        
        return os.path.join(appdata_local, 'rois.json')

    def update_window_title(self):
        """更新窗口标题显示配置文件路径"""
        base_title = "Android设备重定位工具"
        if self.config_path:
            # 显示完整路径
            title = f"{base_title} - {self.config_path}"
        else:
            title = f"{base_title} - 未指定配置文件"
        
        self.setWindowTitle(title)

    def init_ui(self):
        """初始化用户界面"""
        # 初始设置窗口标题（稍后会通过update_window_title更新）
        self.setWindowTitle("Android设备重定位工具")
        self.setGeometry(100, 100, 1000, 700)

        # 创建菜单栏
        self.create_menu_bar()

        # 创建中央控件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 创建布局
        main_layout = QVBoxLayout(central_widget)

        # 创建视频显示控件
        self.video_widget = VideoWidget(self)
        main_layout.addWidget(self.video_widget)

        # 创建状态栏
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        self.update_status("正在连接...")

    def create_menu_bar(self):
        """创建菜单栏"""
        menubar = self.menuBar()
        
        # 创建File菜单
        file_menu = menubar.addMenu('File')
        
        # 添加Load菜单项
        load_action = QAction('Load', self)
        load_action.setShortcut('Ctrl+O')
        load_action.setStatusTip('加载ROI配置文件')
        load_action.triggered.connect(self.load_config_dialog)
        file_menu.addAction(load_action)
        
        # 添加Save菜单项
        save_action = QAction('Save', self)
        save_action.setShortcut('Ctrl+S')
        save_action.setStatusTip('保存ROI配置文件')
        save_action.triggered.connect(self.save_config_dialog)
        file_menu.addAction(save_action)
        
        # 添加Save As菜单项
        save_as_action = QAction('Save As', self)
        save_as_action.setShortcut('Ctrl+Shift+S')
        save_as_action.setStatusTip('另存为ROI配置文件')
        save_as_action.triggered.connect(self.save_as_config_dialog)
        file_menu.addAction(save_as_action)
        
        # 添加分隔符
        file_menu.addSeparator()
        
        # 添加其他功能菜单项
        reconnect_action = QAction('重新连接', self)
        reconnect_action.setShortcut('Ctrl+R')
        reconnect_action.setStatusTip('重新连接设备')
        reconnect_action.triggered.connect(self.manual_reconnect)
        file_menu.addAction(reconnect_action)
        
        clear_action = QAction('清除所有选框', self)
        clear_action.setShortcut('Ctrl+D')
        clear_action.setStatusTip('清除所有ROI选框')
        clear_action.triggered.connect(self.clear_all_rois)
        file_menu.addAction(clear_action)

    def load_config_dialog(self):
        """显示加载配置文件对话框"""
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "加载ROI配置文件",
            "",
            "JSON Files (*.json);;All Files (*)"
        )
        
        if file_path:
            self.config_path = file_path
            self.update_window_title()  # 更新窗口标题
            self.load_config()

    def save_config_dialog(self):
        """显示保存配置文件对话框"""
        if not self.config_path:
            # 如果没有当前配置路径，使用默认路径
            self.config_path = self.get_default_config_path()
            self.update_window_title()  # 更新窗口标题
        
        self.save_config()
        QMessageBox.information(self, "保存成功", f"配置已保存到:\n{self.config_path}")

    def save_as_config_dialog(self):
        """显示另存为配置文件对话框"""
        # 获取默认保存路径
        default_path = self.get_default_config_path()
        
        file_path, _ = QFileDialog.getSaveFileName(
            self,
            "另存为ROI配置文件",
            default_path,
            "JSON Files (*.json);;All Files (*)"
        )
        
        if file_path:
            # 设置新的配置路径
            self.config_path = file_path
            self.update_window_title()  # 更新窗口标题
            # 保存配置
            self.save_config()
            QMessageBox.information(self, "保存成功", f"配置已另存为:\n{self.config_path}")

    def auto_detect_device(self):
        """自动检测连接的Android设备"""
        try:
            result = subprocess.check_output(["adb", "devices"], text=True)
            lines = result.strip().split('\n')
            if len(lines) > 1:
                for line in lines[1:]:
                    if 'device' in line and not line.startswith('*'):
                        return line.split('\t')[0]
            print("未检测到ADB设备，请确保设备已连接")
        except Exception as e:
            print(f"ADB设备检测失败: {e}")
        return None

    def load_config(self):
        """加载配置文件"""
        if self.config_path and os.path.exists(self.config_path):
            try:
                with open(self.config_path, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                self.normalized_rois = config.get('rois', [])
                print(f"已加载配置文件: {len(self.normalized_rois)} 个选框")
                self.update_status(f"已加载配置 - {len(self.normalized_rois)} 个选框")
                # 更新显示
                self.denormalize_rois()
                self.video_widget.update()
            except Exception as e:
                print(f"配置文件加载失败: {e}")
                QMessageBox.warning(self, "加载失败", f"配置文件加载失败:\n{e}")
        else:
            # 如果配置文件不存在，显示默认状态
            self.update_status("配置文件不存在，将创建新配置")

    def save_config(self):
        """保存配置文件"""
        if self.config_path:
            try:
                config = {'rois': self.normalized_rois}
                config_dir = os.path.dirname(self.config_path)
                if config_dir and not os.path.exists(config_dir):
                    os.makedirs(config_dir, exist_ok=True)

                with open(self.config_path, 'w', encoding='utf-8') as f:
                    json.dump(config, f, indent=2, ensure_ascii=False)
                print(f"配置已保存到: {self.config_path}")
                self.update_status(f"配置已保存 - {len(self.normalized_rois)} 个选框")
            except Exception as e:
                print(f"配置文件保存失败: {e}")
                QMessageBox.warning(self, "保存失败", f"配置文件保存失败:\n{e}")

    def normalize_rois(self):
        """将显示ROI坐标归一化"""
        if not self.frame_width or not self.frame_height:
            return

        self.normalized_rois = []
        video_rect = self.video_widget.video_rect
        scale_factor = self.video_widget.scale_factor

        for roi_rect in self.display_rois:
            # 转换为视频坐标系
            video_x = (roi_rect.x() - video_rect.x()) / scale_factor
            video_y = (roi_rect.y() - video_rect.y()) / scale_factor
            video_w = roi_rect.width() / scale_factor
            video_h = roi_rect.height() / scale_factor

            # 归一化
            norm_roi = {
                'x': video_x / self.frame_width,
                'y': video_y / self.frame_height,
                'w': video_w / self.frame_width,
                'h': video_h / self.frame_height
            }
            self.normalized_rois.append(norm_roi)

    def denormalize_rois(self):
        """将归一化ROI转换为显示坐标"""
        if not self.frame_width or not self.frame_height:
            return

        self.display_rois = []
        video_rect = self.video_widget.video_rect
        scale_factor = self.video_widget.scale_factor

        for norm_roi in self.normalized_rois:
            # 反归一化到视频坐标
            video_x = norm_roi['x'] * self.frame_width
            video_y = norm_roi['y'] * self.frame_height
            video_w = norm_roi['w'] * self.frame_width
            video_h = norm_roi['h'] * self.frame_height

            # 转换为显示坐标
            display_x = int(video_x * scale_factor + video_rect.x())
            display_y = int(video_y * scale_factor + video_rect.y())
            display_w = int(video_w * scale_factor)
            display_h = int(video_h * scale_factor)

            roi_rect = QRect(display_x, display_y, display_w, display_h)
            self.display_rois.append(roi_rect)

    def add_roi(self, rect):
        """添加新的ROI"""
        self.display_rois.append(rect)
        self.normalize_rois()
        self.update_status(f"已添加选框 - 总数: {len(self.display_rois)}")

    def delete_roi_at_position(self, pos):
        """删除指定位置的ROI"""
        for i, roi_rect in enumerate(self.display_rois):
            if roi_rect.contains(pos):
                self.display_rois.pop(i)
                self.normalize_rois()
                self.update_status(f"已删除选框 - 剩余: {len(self.display_rois)}")
                self.video_widget.update()
                break

    def clear_all_rois(self):
        """清除所有ROI"""
        self.display_rois.clear()
        self.normalized_rois.clear()
        self.update_status("已清除所有选框")
        self.video_widget.update()

    def manual_reconnect(self):
        """手动重连"""
        self.connected = False
        self.reconnect_attempts = 0
        self.connect_client()

    def update_status(self, message):
        """更新状态栏"""
        self.status_bar.showMessage(message)

    def frame_callback(self, client, frame):
        """scrcpy帧回调函数"""
        if frame is None:
            return

        try:
            self.connected = True
            self.reconnect_attempts = 0

            # 更新帧尺寸
            if self.frame_width != frame.shape[1] or self.frame_height != frame.shape[0]:
                original_height, original_width = frame.shape[:2]

                # 按目标宽高比调整
                if original_width / original_height > self.aspect_ratio:
                    self.frame_height = original_height
                    self.frame_width = int(original_height * self.aspect_ratio)
                else:
                    self.frame_width = original_width
                    self.frame_height = int(original_width / self.aspect_ratio)

            with self.frame_lock:
                self.latest_frame = frame.copy()

        except Exception as e:
            print(f"帧处理错误: {e}")

    def connect_client(self):
        """连接scrcpy客户端"""
        try:
            if self.client:
                try:
                    self.client.stop()
                except:
                    pass
                self.client = None

            if not self.device_ip:
                self.update_status("未找到设备")
                return False

            self.update_status(f"正在连接设备: {self.device_ip}")

            self.client = Client(
                device=self.device_ip,
                max_fps=30,
                bitrate=4000000,
                max_size=self.max_size,
                block_frame=True,
                skip_same_frame=False
            )

            self.client.on_frame(self.frame_callback)
            self.client.start(threaded=True)

            # 等待连接
            QTimer.singleShot(3000, lambda: self.update_status("已连接") if self.connected else self.update_status("连接超时"))

            return True

        except Exception as e:
            self.update_status(f"连接失败: {e}")
            return False

    def update_display(self):
        """更新视频显示"""
        with self.frame_lock:
            if self.latest_frame is None:
                return

            frame = self.latest_frame.copy()

        try:
            # 按目标宽高比裁剪
            original_height, original_width = frame.shape[:2]

            if original_width / original_height > self.aspect_ratio:
                # 裁剪宽度
                crop_width = int(original_height * self.aspect_ratio)
                crop_height = original_height
                crop_x = (original_width - crop_width) // 2
                crop_y = 0
            else:
                # 裁剪高度
                crop_width = original_width
                crop_height = int(original_width / self.aspect_ratio)
                crop_x = 0
                crop_y = (original_height - crop_height) // 2

            # 裁剪帧
            cropped_frame = frame[crop_y:crop_y+crop_height, crop_x:crop_x+crop_width]

            # 转换为QImage
            height, width, channel = cropped_frame.shape
            bytes_per_line = 3 * width
            q_image = QImage(cropped_frame.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()

            # 计算显示参数
            widget_size = self.video_widget.size()
            image_size = q_image.size()

            # 保持宽高比缩放
            scaled_size = image_size.scaled(widget_size, Qt.KeepAspectRatio)
            scaled_pixmap = QPixmap.fromImage(q_image).scaled(scaled_size, Qt.KeepAspectRatio, Qt.SmoothTransformation)

            # 计算视频在控件中的位置
            x_offset = (widget_size.width() - scaled_size.width()) // 2
            y_offset = (widget_size.height() - scaled_size.height()) // 2
            video_rect = QRect(x_offset, y_offset, scaled_size.width(), scaled_size.height())
            scale_factor = scaled_size.width() / width

            # 更新视频控件参数
            self.video_widget.set_video_rect(video_rect, scale_factor)

            # 更新ROI显示坐标
            self.denormalize_rois()

            # 显示图像
            self.video_widget.setPixmap(scaled_pixmap)

        except Exception as e:
            print(f"显示更新错误: {e}")

    def closeEvent(self, event):
        """窗口关闭事件"""
        if self.client:
            try:
                self.client.stop()
            except:
                pass
        event.accept()

def main():
    parser = argparse.ArgumentParser(description='Android设备重定位工具')
    parser.add_argument('-c', '--config', type=str, 
                        help='选框配置文件路径 (默认: AppData/Local/AndroidRelocator/rois.json)')
    parser.add_argument('-d', '--device', type=str, help='指定设备IP地址')
    parser.add_argument('--aspect-ratio', type=float, default=1.8526,
                        help='视频显示宽高比 (默认: 1.8526)')
    parser.add_argument('-m', '--max-size', type=int, default=1778,
                        help='视频最大尺寸 (默认: 1778)')

    args = parser.parse_args()

    app = QApplication(sys.argv)

    window = AndroidRelocatorGUI(
        device_ip=args.device,
        config_path=args.config,  # 如果为None，会使用默认AppData路径
        aspect_ratio=args.aspect_ratio,
        max_size=args.max_size
    )

    window.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()
